package com.example.demo.service.aiService;

import com.example.demo.demo.invoke.DialogueState;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.stereotype.Service;

@Service
public class StateExtractionService {
    private final ChatClient chatClient;
    private final ObjectMapper objectMapper;
    private static final String STATE_EXTRACTION_PROMPT = """
            You are a Natural Language Understanding engine. Your task is to analyze the dialogue context and update the dialogue state based on the user's latest message.
            Your final output MUST BE a single, valid JSON object representing the **fully updated** dialogue state, containing all information from both the current state and the new message. The JSON must match the DialogueState class structure.

            ### Current Dialogue State (JSON) ###
            {currentStateJson}

            ### User's Latest Message ###
            "{latestQuery}"

            ### Your Task ###
            Generate the JSON for the new, updated dialogue state.
            """;

    public StateExtractionService(ChatClient.Builder chatClientBuilder, ObjectMapper objectMapper) {
        this.chatClient = chatClientBuilder.build();
        this.objectMapper = objectMapper;
    }

    public DialogueState updateState(DialogueState currentState, String latestQuery) {
        String currentStateJson = convertStateToJson(currentState);
        return this.chatClient.prompt()
                .user(userSpec -> userSpec
                        .text(STATE_EXTRACTION_PROMPT)
                        .param("currentStateJson", currentStateJson)
                        .param("latestQuery", latestQuery))
                .call()
                .entity(DialogueState.class);
    }

    private String convertStateToJson(DialogueState state) {
        try {
            return objectMapper.writeValueAsString(state);
        } catch (JsonProcessingException e) {
            throw new RuntimeException("Error serializing DialogueState to JSON", e);
        }
    }
}
